"""
使用fasttext获取词向量
"""
import fasttext
import chatbot.config as config


def build_model(by_word=True):
    data_path = "corpus/dnn/sort/sim_q_cuted_words.txt" if by_word else "corpus/dnn/sort/sim_q_cuted.txt"
    model = fasttext.train_unsupervised(data_path, wordNgrams=1, epoch=20, minCount=2)
    model.save_model(config.recall_fasttext_by_word_save_path if by_word else config.recall_fasttext_save_path)


def get_model(by_word=True):
    if by_word:
        return fasttext.load_model(config.recall_fasttext_by_word_save_path)
    else:
        return fasttext.load_model(config.recall_fasttext_save_path)


class FastTextVectorizer():
    def __init__(self, by_word=True):
        self.model = get_model(by_word)

    def transform(self, sentences):
        """
        :param sentences: [sentence,sentence.....]
        :return:
        """
        results = [self.model.get_sentence_vector(i) for i in sentences]
        return results

    def fit_transform(self, sentences):
        return self.transform(sentences)
